function res = traintest(stego, cover, method, varargin)
% TRAINTEST train models use various methods.
% USAGE:
%   RES = TRAINTEST(STEGO, COVER, ...)
%
% INPUT:
%   STEGO  : a L*N matrix for features of stego images.
%   COVER  : a M*N matrix for features of cover images.
%
% OUTPUT:
%

%% Parse input, set default values
param = parse_args(method, varargin{:});
if ischar(cover)
    if param.verbose
        fprintf('loading cover from file %s.\n', cover);
    end
    load(cover,'-mat','res');
    cover = res;
    clear res;
end
if ischar(stego)
    if param.verbose
        fprintf('loading stego from file %s.\n', stego);
    end
    load(stego,'-mat','res');
    cover = res;
    clear res;
end

[L, N1] = size(stego);
[M, N2] = size(cover);
if (N1 ~= N2)
    error('dimension of stego and cover feature must match.');
end

%% prepare return value, store param
res = struct('method',method, 'STEGO_SAMPLE_NUM', L, ...
    'COVER_SAMPLE_NUM', M, 'FEATURE_DIM', N1);

%% select parameters
switch lower(method)
    case 'libsvm'
        if strcmpi(param.c, 'automatic') || strcmpi(param.g, 'automatic') || isnan(param.c) || isnan(param.g)
            cg_option = [param.cg_opt, '-t ',num2str(param.kernel_type)];
            if param.verbose
                fprintf('finding best c and g\n');
            else
                cg_option = ['-q ',cg_option];
            end
        end        
        [trainset, trainlbl] = splitset(stego, cover, param);
        t1 = cputime;
        [acc, c, g, cgmat] = SVMcg(trainlbl, trainset, cg_option);
        res.CG_TIME = cputime - t1;
        
        [res.CG_ACC, param.c, param.g, res.CG_MAT] = deal(acc, c, g, cgmat);
        if param.verbose
            disp(cvmat);
        end
        fprintf('c=%g, g=%g\n', param.c, param.g);
        clear trainlbl trainset cg_option
    case 'ensemble'
        if param.verbose
            fprintf('finding best d_sub and l\n');
        end
        ak = zeros(param.arg_loop, 1);
        al = ak;
        ae = ak;
        
        for i = 1:param.arg_loop
            if param.verbose
                fprintf('round %d/%d\n', i, param.arg_loop);
            else
                fprintf('.');
            end
            [trainset, trainlbl] = splitset(stego, cover, param);
            pstego = (trainlbl == 1);
            pcover = (trainlbl == 0);

            t1 = cputime;
            model = ensemble_wrap(trainset(pstego,:), trainset(pcover, :), 'same', 'automatic', 'automatic', 0.5, param.verbose);
            train_time(i) = cputime - t1;
            
            ae(i) = model.testing_error;
            ak(i) = model.optimal.k;
            al(i) = model.optimal.L;
        end
        
        param.d_sub = round(mean(ak));
        param.L = round(mean(al));
        res.KL_ERR = mean(ae);

    otherwise
        % SKIP THIS SECTION
end

res.param = param;
if param.verbose
    disp(param);
end
[res.MEAN_ERR, res.MEAN_FA, res.MEAN_MD, res.MIN_ERR, res.MAX_ERR] = deal(nan);

if param.loop <= 0
    error('trainModel:invalidInput','Loop must be >0.');
end
errors_fa = nan(param.loop, 1); %false alarm 0 said to be 1
errors_md = nan(param.loop, 1); %miss detect 1 said to be 0
%accs = nan(param.loop, 3);
errors = nan(param.loop, 1);
train_time = nan(param.loop, 1);
test_time = nan(param.loop, 1);
es = 0;

%% train and test in loops
cloop = 0;
for i = 1:param.loop
    if param.verbose
        fprintf('Split training and testing sets ');
    end
    [trainset, trainlbl, testset, testlbl] = splitset(stego, cover, param);
    if param.verbose
        fprintf('done\n');
        fprintf('trn/tst = %d/%d\n',length(trainlbl), length(testlbl));
    end
                
    switch lower(method)
        case 'libsvm'
            train_option = [param.train_opt, '-t ',num2str(param.kernel_type),' -c ',num2str(param.c), ' -g ',num2str(param.g)];
            test_option = [param.predict_opt];
            if ~param.verbose
                train_option = ['-q ',train_option];
                test_option = ['-q ',test_option];
            end
            t1 = cputime;
            model = libsvmtrain(trainlbl, trainset, train_option)
            train_time(i) = cputime - t1;
            
            t1 = cputime;
            [tr, acc, ~] = libsvmpredict(testlbl, testset, model, test_option);
            test_time(i) = cputime - t1;
            
        case 'lda'
        case 'kda'
        case 'ensemble'
            pstego = (trainlbl == 1);
            pcover = (trainlbl == 0);
            t1 = cputime;
            model = ensemble_train_wrap(trainset(pstego,:), trainset(pcover, :), param.d_sub, param.L, 'same', param.verbose);
            train_time(i) = cputime - t1;
            
            t1 = cputime;
            tr = ensemble_predict(testMat, model);
            test_time(i) = cputime - t1;
            
        otherwise
            error('Unknown method');
    end
    pfa = sum((tr == 1) & (testlbl == 0)) / sum(testlbl == 0);
    pmd = sum((tr == 0) & (testlbl == 1)) / sum(testlbl == 1);
    err = (pmd + pfa) / 2;
    errors_fa(i) = pfa;
    errors_md(i) = pmd;
    errors(i) = err;
    es = es + err;
    if param.verbose
        fprintf('accuracy = %0.2f%%, average = %0.2f%%\n',(1-err)*100,(1 - es / i)*100);
    end
    cloop = cloop + 1;
end

%% collect stat results and output
res.MIN_ERR = min(errors);
res.MAX_ERR = max(errors);
errors(isnan(errors)) = 0;
errors_fa(isnan(errors_fa)) = 0;
errors_md(isnan(errors_md)) = 0;
[res.MEAN_ERR, res.MEAN_FA, res.MEAN_MD] = deal(sum(errors)/param.loop, sum(errors_fa)/param.loop, sum(errors_md) / param.loop);
res.ERRORS_FA = errors_fa;
res.ERRORS_MD = errors_md;
res.ERRORS = errors;

if param.verbose
    disp(res)
    %fprintf('mean(accuracy) = %0.2f%%\n',(1-res.MEAN_ERR)*100);
end
if ischar(param.output)
    % save model
    save(param.output,'model', 'param', 'res');
end
end

function exportSVM(fname, lbl, inst)
fid = fopen(fname, 'w');
if (fid == -1)
    error('Cannot open file %s for write', fname);
end
[M,N]= size(inst);
for i = 1:M
    fprintf(fid,'%d',lbl(i));
    for j = 1:N
        fprintf(fid,' %d:%0.4f',j,inst(i,j));
    end
    fprintf(fid,'\n');
end
fclose(fid);
end

function param = parse_args(method, varargin)
param = default_args(method);
n = length(varargin);
i = 1;
while (i <= n)
    arg = varargin{i};
    if ~ischar(arg)
        error('key must be a string.');
    end
    if (i + 1 > n)
        error('arguments should be in pairs, e.g., key1, value1, key2, value2, ...');
    end
    val = varargin{i+1};
    i = i + 2;
    param.(arg) = val;
end
end

function param = default_args(method)
param = struct('verbose', 1, 'pair', 1, 'loop', 12, 'output',0);
switch lower(method)
    case 'libsvm'
        param.kernel_type = 0;
        param.cg_opt = '-h 1 -m 512';
        param.train_opt = '-h 1 -m 512';
        param.predict_opt = '';
        param.c = 'automatic';
        param.g = 'automatic';
        param.sample_cg = 0;
    case 'ensemble'
        param.L = 'automatic';
        param.d_sub = 'automatic';
        param.arg_loop = 5;
end
end

function [trainset, trainlbl, testset, testlbl] = splitset(stego, cover, param)
% split training sets and testing sets
[s1, h1] = size(stego);
[s2, h2] = size(cover);
s = s1+s2;
if h1 ~= h2
    error('dimension of stego and cover feature must match.');
end

% ssr - start of stego images in training sets = 1
% nsr - number of stego images chosen in training sets
% sse - start of stego images in testing sets
% nse - number of stego images chosen in testing sets
% scr - start of cover images in training sets = 1
% ncr - number of cover images chosen in training sets
% sce - start of cover images in testing sets
% nce - number of cover images chosen in testing sets

ssr = 1;
scr = 1;
mode = 1;
if isfield(param, 'ratio')
    if ~param.pair
        % only choose a selected ratio
        nsr = round(s1 * ratio);
        ncr = round(s2 * ratio);
        nse = s1 - nsr;
        nce = s2 - ncr;
    else
        % only balance stego / cover image in training sets
        r = round(s * param.ratio / 2)*2;
        nsr = r / 2;
        ncr = r - nsr;
        nse = s1 - nsr;
        nce = s2 - ncr;
    end
else
    % only ntrain / ntest specified
    if param.pair
        nsr = round(param.ntrain / 2);
        ncr = nsr;
        nse = round(param.ntest / 2);
        nce = nse;
    else
        r = param.ntrain;
        e = param.ntest;
        mode = 2;
    end
end
% mode = 1 position is decided
% mode = 2 position is not decided
if (mode == 1)
    sse = nsr + ssr;
    sce = ncr + scr;
    if (ncr <= 0) || (nsr <= 0) || (nce <= 0) || (nse <= 0)
        error('size of training sets or testing sets must > 0.');
    end
    if (scr + ncr > sce) || (sce + nce > s2+1) || (ssr + nsr > sse) || (sse + nse > s1+1)
        error('size of training sets or testing sets must <= number of samples.');
    end
    trainlbl = [ones(nsr,1); zeros(ncr,1)];
    testlbl = [ones(nse,1); zeros(nce,1)];
    trainset = [stego(ssr:ssr+nsr-1,:); cover(scr:scr+ncr-1,:)];
    testset = [stego(sse:sse+nse-1,:); cover(sce:sce+nce-1,:)];
elseif (mode == 2)
    lbl = [ones(s1,1);zeros(s2,1)];
    data = [stego;cover];
    p = randperm(s);
    trainlbl = lbl(p(1:r),:);
    testlbl = lbl(p(r+1:r+e),:);
    trainset = data(p(1:r),:);
    testset = data(p(r+1:r+e),:);
end
end